from openai import OpenAI
from pymilvus import MilvusClient

from .embeddings import get_embedding_function
from ..vanna.milvus import Milvus_VectorStore
from ..vanna.openai import OpenAI_Chat


class DQuestionMilvus(Milvus_VectorStore, OpenAI_Chat):
    def __init__(self, config=None):
        client = None
        if config is None:
            milvus_client = MilvusClient(uri="http://122.51.212.243:19530")
            client = OpenAI(api_key="sk-aeb8d69039b14320b0fe58cb8285d8b1",
                            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1")

            config = {"milvus_client": milvus_client,
                      "embedding_function": get_embedding_function(),
                      "dialect": "SQLLite",
                      "language": "Chinese",
                      "model": "deepseek-v3",
                      "dim": 1024,
                      "use_hybrid_search":True
                      }
        Milvus_VectorStore.__init__(self, config=config)
        OpenAI_Chat.__init__(self, client=client, config=config)


class SQLiteDatabaseRAG(DQuestionMilvus):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.kwargs = kwargs
        self.connect_to_sqlite('https://vanna.ai/Chinook.sqlite')

    async def load_data(self):
        pass

    async def create_index(self):
        # 获取数据库中的DDL信息
        df_ddl = self.run_sql("SELECT type,sql FROM sqlite_master WHERE sql is not null")
        current_ddls = set(df_ddl['sql'].to_list())
        # 获取已有的训练数据
        existing_training_data = self.get_training_data()
        existing_ddls = set()
        # 提取已有训练数据中的DDL内容
        for _, row in existing_training_data.iterrows():
            if row['content'] is not None:
                existing_ddls.add(row['content'])
        # 找出需要新增的DDL
        new_ddls = current_ddls - existing_ddls
        # 找出需要删除的DDL
        obsolete_ddls = existing_ddls - current_ddls
        # 删除过时的训练数据
        for ddl in obsolete_ddls:
            # 查找包含该DDL的训练数据ID
            matching_rows = existing_training_data[existing_training_data['content'] == ddl]
            for _, row in matching_rows.iterrows():
                self.remove_training_data(row['id'])
        # 添加新的DDL训练数据
        for ddl in new_ddls:
            self.train(ddl=ddl)


class MySQLDatabaseRAG(DQuestionMilvus):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.kwargs = kwargs
        self.connect_to_mysql(host="122.51.212.243",
                              port=3306,
                              user="root",
                              password="123456",
                              dbname="testPerformance"
                              )

    async def load_data(self):
        pass

    async def create_index(self):
        df_ddl = self.run_sql(
            f"SELECT * FROM INFORMATION_SCHEMA.COLUMNS where table_schema = 'testPerformance'")
        plan = self.get_training_plan_mysql(df_ddl)
        self.train(plan=plan)
